Code Review: PR #524 — fix: persist side-effect columns to row buffer in async engine

Summary

This PR fixes a bug where the async engine's _run_cell and _run_batch methods silently dropped side-effect columns (e.g., __trace, __reasoning_content) produced by generators but not tracked in _instance_to_columns. Downstream columns referencing these values would fail with missing-column errors.

The fix adds a second write pass in both _run_cell and _run_batch that persists any new keys from the generator result that were not present in the input row data. A unit test is included for the cell-by-cell code path.

Files changed: 2 (86 additions, 1 deletion)

• packages/data-designer-engine/src/data_designer/engine/dataset_builders/async_scheduler.py
• packages/data-designer-engine/tests/engine/dataset_builders/test_async_scheduler.py

CI status: Lint passes. Most test/E2E jobs still pending at review time. Config tests all pass.

Findings

1. _run_from_scratch not updated — potential silent drop (Medium)

File: async_scheduler.py:757-775

The _run_from_scratch method has the same pattern as _run_cell and _run_batch — it only writes output_cols to the buffer. If a from-scratch generator ever produces side-effect columns, they will be silently dropped, exactly the same bug being fixed here for the other two paths.

This may be intentional (from-scratch generators might not produce side-effects today), but it creates an asymmetry across the three execution paths. At minimum, a comment explaining why _run_from_scratch doesn't need the same treatment would prevent a future reader from filing the same bug again.

2. No test coverage for _run_batch side-effect handling (Medium)

The test (test_scheduler_side_effect_columns_written_to_buffer) only exercises cell-by-cell generation via MockCellGeneratorWithSideEffect. The parallel change to _run_batch (lines 834-840) has no corresponding test. Since the batch path has different mechanics (iterating over a DataFrame vs. a dict, tracking existing_cols from the batch input), a dedicated test would reduce regression risk.

3. Side-effect column collision semantics are implicit (Low)

In _run_cell, the filter col not in row_data means that if two sequential generators both produce the same side-effect column name, the second generator's value is silently dropped (the first writer wins). This is likely fine given the naming convention (column__reasoning_content), but a brief comment noting the first-writer-wins behavior would help future maintainers understand the design choice.

4. Minor: batch_df is not None check may be unnecessary (Nit)

File: async_scheduler.py:837

existing_cols = set(batch_df.columns) if batch_df is not None else set()

batch_df is assigned from self._buffer_manager.get_dataframe(...) just above, and we are inside the if self._buffer_manager is not None guard. If get_dataframe can never return None, this defensive check is dead code. If it can, this is fine. Minor either way.

Positives

• Focused, minimal fix. Only the necessary lines are changed, no unrelated refactoring.
• Good comments. The inline comments explaining what side-effect columns are and why they need persisting add valuable context.
• Clean test design. MockCellGeneratorWithSideEffect is a clear, well-scoped test helper. The test verifies both presence and content of the side-effect column across all rows.
• Correct set operations. Using result.keys() - written for the cell path is idiomatic and efficient.

Verdict

Approve with suggestions. The core fix is correct and well-implemented. The two medium findings (missing _run_from_scratch symmetry and missing batch test) are worth addressing to complete the fix, but are not blockers for merge. The low/nit items are optional improvements.

Greptile Summary

This PR fixes a silent data-loss bug in the async scheduler where side-effect columns produced by generators (e.g. __trace, __reasoning_content) were never written back to the RowGroupBufferManager, causing downstream columns that reference them to fail with missing-column errors.

The fix introduces a separate _instance_to_write_columns dict (a superset of _instance_to_columns) that includes side-effect columns alongside primary ones, so all three write-back paths (_run_from_scratch, _run_cell, _run_batch) now persist them. _instance_to_columns is intentionally left unchanged for completion tracking and dispatch dedup.

Important Files Changed

sequenceDiagram
    participant DB as DatasetBuilder
    participant AS as AsyncTaskScheduler
    participant Gen as ColumnGenerator
    participant BM as RowGroupBufferManager

    DB->>DB: build side_effect_map
    DB->>AS: AsyncTaskScheduler(..., side_effect_map)
    AS->>AS: build _instance_to_write_columns

    loop Cell-by-cell task for answer
        AS->>BM: get_row(rg, ri)
        AS->>Gen: agenerate(row_data)
        Gen-->>AS: result with answer + answer__reasoning_content
        AS->>BM: update_cell answer
        AS->>BM: update_cell answer__reasoning_content
    end

    loop Cell-by-cell task for judge
        AS->>BM: get_row includes answer__reasoning_content
        AS->>Gen: agenerate(row_data)
    end

_{Reviews (2): Last reviewed commit: "try a better fix" | Re-trigger Greptile}

Hey @nabinchha - heads up that #509 addresses the same root cause and also fixes a related bug in ExecutionGraph.set_side_effect (first-writer-wins semantics to prevent false DAGCircularDependencyError when multiple stages declare the same side-effect). The scheduler fix in #509 also avoids the extra wiring in dataset_builder.py by reading side-effect columns directly from generator configs inside AsyncTaskScheduler.__init__.

Want to sync up so we're not duplicating effort? Happy to merge the approaches - the _instance_to_write_columns split is the right call either way.

@andreatgretel you are right! This is a duplicate of #509. Will close.